package com.gitee.wsl.flow.transform.split

import com.gitee.wsl.flow.transform.TransformationContext
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.channels.BufferOverflow
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.flow.flatMapMerge
import kotlinx.coroutines.flow.onStart
import kotlinx.coroutines.flow.receiveAsFlow


/**
 * Transforms a [Flow] of [Input] to a [Flow] of [Output] by splitting the original into [Flow]s
 * of type [Selector]. Each independent [Flow] of the [Selector] type can then be transformed
 * into a [Flow] of [Output].
 *
 * [typeSelector]: The mapping to the type the [Input] [Flow] should be split into
 *
 * [keySelector]: The mapping to the [Selector] to the key used to identify it. This is useful
 * for nested class hierarchies. By default each distinct type will be split out, but if you want
 * to treat certain subtypes as one type, this lets you do that.
 * [transform]: a function for mapping independent [Flow]s of [Selector] to [Flow]s of [Output]
 */
@OptIn(ExperimentalCoroutinesApi::class)
fun <Input : Any, Selector : Any, Output : Any> Flow<Input>.splitByType(
    capacity: Int,
    onBufferOverflow: BufferOverflow,
    typeSelector: (Input) -> Selector,
    keySelector: (Selector) -> String = Any::defaultKeySelector,
    // Ergonomic hack to simulate multiple receivers
    transform: suspend TransformationContext<Selector>.() -> Flow<Output>
): Flow<Output> =
    channelFlow mutationFlow@{
        val keysToFlowHolders = mutableMapOf<String, FlowHolder<Selector>>()
        this@splitByType
            .collect { item ->
                val selected = typeSelector(item)
                val flowKey = keySelector(selected)
                when (val existingHolder = keysToFlowHolders[flowKey]) {
                    null -> {
                        val holder = FlowHolder(
                            capacity = capacity,
                            onBufferOverflow = onBufferOverflow,
                            firstEmission = selected
                        )
                        keysToFlowHolders[flowKey] = holder
                        val context = TransformationContext(selected, holder.exposedFlow)
                        val mutationFlow = transform(context)
                        channel.send(mutationFlow)
                    }

                    else -> {
                        existingHolder.internalSharedFlow.send(selected)
                    }
                }
            }
    }
        .flatMapMerge(
            concurrency = Int.MAX_VALUE,
            transform = { it }
        )

/**
 * Container for representing a [Flow] of a subtype of [Action] that has been split out from
 * a [Flow] of [Action]
 */
private data class FlowHolder<Action>(
    val capacity: Int,
    val onBufferOverflow: BufferOverflow,
    val firstEmission: Action,
) {
    val internalSharedFlow: Channel<Action> = Channel(
        capacity = capacity,
        onBufferOverflow = onBufferOverflow,
    )
    val exposedFlow: Flow<Action> = internalSharedFlow
        .receiveAsFlow()
        .onStart { emit(firstEmission) }
}

private fun Any.defaultKeySelector(): String = this::class.simpleName
    ?: throw IllegalArgumentException(
        "Only well defined classes can be split or specify a different key selector"
    )